from torch import nn

class NeuralNetwork(nn.Module):

    def __init__(self):
        super().__init__()
        self.lstm = nn.LSTM(input_size=33, hidden_size=33, num_layers=1, batch_first=True)
        self.linear_out = nn.Linear(in_features=33, out_features=33)

    def forward(self, input_data):
        lstm_out, _ = self.lstm(input_data)
        return self.linear_out(lstm_out[:, -1, :])

if __name__ == '__main__':
    from dataset import iter_data
    import torch
    from torch.nn import Softmax

    sm = Softmax(dim=1)

    with torch.no_grad():
        model = NeuralNetwork()
        print(model)
        model.eval()
        for x, y in iter_data(2):
            print('model input:', x)
            print('model output:', model(x))
            print('y:', y)
            break
